from langchain_core.messages import AIMessage

from langgraph_coder.src.agent.chain.generate_chain import GenerateChain
from langgraph_coder.src.agent.model.graph_state import GraphState


class GenerateNode:
    """
    生成代码结点
    """
    def __init__(self, code_gen_chain: GenerateChain):
        self.code_gen_chain = code_gen_chain

    def __call__(self, state: GraphState):
        """
        生成代码
        :param state: 图结构
        :return:
        """
        print("---代码生成中---")

        # 拿到相关参数
        messages = state["messages"]
        iterations = state["iterations"]
        error = state["error"]

        # 出现错误返回生成阶段，并反馈相关错误
        if error == "yes":
            messages += [
                (
                    "user",
                    "请再试一次。调用代码工具以结构化输出，包括prefix、imports和code block:",
                )
            ]

        # 利用之前构建的链类（基于大模型实现）来获取代码解决方案
        code_solution = self.code_gen_chain.get_solution(messages)

        # 维护这个历史对话
        messages += [AIMessage(content=f"{code_solution['prefix']} \n Imports: {code_solution['imports']} \n Code: {code_solution['code']}")]

        # 添加迭代次数
        iterations = iterations + 1

        # 返回图结构
        state["messages"] =  messages
        state["iterations"] = iterations
        state["generation"] = code_solution
        return state

